#------------------------------------------------------------------------------
# plot marginal effects for Figure 1, 2, Appendix Figure S1.
#==============================================================================
library(tidyverse)
library(data.table)
library(hrbrthemes)
library(extrafont)
library(viridis)
library(ggpubr)
library(ggsci)

# main effects
if (!dir.exists(figure_path)){
	dir.create(figure_path)
}

# load code that is necessary for plotting Figure 1,2,3.
source(file.path(code_path,'9_function_figure.R'))

# read all files 
coef_combine = list(); n = 1
	coef_combine[[n]] = read_output(log_path=log_path,sameness='raw',mm=1);n = n + 1
  coef_combine[[n]] = read_output(log_path=log_path,sameness='sameness',mm=3);n = n + 1
	coef_combine[[n]] = read_output_mi(log_path=log_path,sameness='raw',mm=2);n = n + 1
	coef_combine[[n]] = read_output_mi(log_path=log_path,sameness='sameness',mm=4);n = n + 1
	
coef_combine = rbindlist(coef_combine,fill=TRUE)

#------------------------------------------------------------------------------
# main effets (raw)
#==============================================================================
coef_for_plot = coef_combine[sameness=='raw'&grepl('vs|Range',mtype),]
coef_for_plot = coef_for_plot[grep('vs 2|vs 3|vs 4',mtype,invert=TRUE),]
coef_for_plot = coef_for_plot[at_type %in% c('Other factors','Religious Congregation')== FALSE,]

coef_for_plot = rbind(coef_for_plot[method=='logit',],
	coef_for_plot[method=='mi' & category %in% c('UnEmpl','PhysProb','prop_UnEmpl','prop_PhysProb'),])

coef_for_plot[, group := ifelse(grepl('\\(vs ',category_name), 'individual','contextual')]
coef_for_plot[, group := factor(group, levels=c('individual','contextual'))]
coef_for_plot[,table(category_name)]

coef_for_plot[method=='logit',sig :=ifelse(p < 0.05, 'sig.','insig.')]
coef_for_plot[method=='mi',sig := ifelse(ll >0 | ul < 0, 'sig.','insig.')]

type_list = unique(coef_for_plot[,at_type])

coef_for_plot[,labels :=  round(change*100000)]

coef_for_plot[grep('Female',category_name), category_order := 1]
coef_for_plot[grep('Age',category_name), category_order := 2]
coef_for_plot[grep('Married|Widowed|Divorced|Separated|Never-married',category_name), category_order := 3]
coef_for_plot[grep('White|Black|Indian|Asian|Hispanic',category_name), category_order := 4]
coef_for_plot[grep('Native',category_name), category_order := 5]
coef_for_plot[grep('Unemployed',category_name), category_order := 6]
coef_for_plot[grep('Problem',category_name), category_order := 7]

coef_for_plot[change < 0 & sig == 'sig.',colors := 'neg_sig']
coef_for_plot[change < 0 & sig == 'insig.',colors := 'neg_insig']
coef_for_plot[change > 0 & sig == 'sig.',colors := 'pos_sig']
coef_for_plot[change > 0 & sig == 'insig.',colors := 'pos_insig']

coef_for_plot[,category_name:=gsub('\\[range\\]','',category_name)]
coef_for_plot[category_name =='Native (vson-native)',category_name:='Native (vs non-native)']

file_name = 'figure1_lolly_pop_main_effects_county.pdf'

p1 = ggplot(coef_for_plot[group=='individual',], aes(x=reorder(category_name,-category_order),
	y=change*100000,yend=ll*100000,ymax=ul*100000,label=round(change*100000,1),lty=sig,color=colors)) + 
  geom_point(stat='identity', size=6)  +
  geom_hline(yintercept=0,lty=2,lwd=0.1)+
  geom_segment(aes(y = ll*100000, 
                   x = category_name, 
                   yend = ul*100000, 
                   xend = category_name), 
               color = "black") +
  scale_color_manual(values=c('neg_sig'='blue','neg_insig'='skyblue','pos_sig'='red','pos_insig'='pink'))+
  scale_linetype_manual(values=c('insig.'=2,'sig.'=1))+
  geom_text(color="white", size=2) +
  labs(subtitle='Individual',
       y='AME on suicide rates (per 100k)', x='') + 
  theme_pubr()+
  coord_flip() +
  theme(
    axis.text.y=element_text(size=8),
    axis.text.x=element_text(size=8),
    axis.title=element_text(size=8),
    plot.subtitle = element_text(hjust = 0.5))+
  guides(lty=FALSE,color=FALSE)

p2 = ggplot(coef_for_plot[group=='contextual',], aes(x=reorder(category_name,-category_order),
  y=change*100000,yend=ll*100000,ymax=ul*100000,label=round(change*100000,1),lty=sig,color=colors)) + 
  geom_point(stat='identity', size=6)  +
  geom_hline(yintercept=0,lty=2,lwd=0.1)+
  geom_segment(aes(y = ll*100000, 
                   x = category_name, 
                   yend = ul*100000, 
                   xend = category_name), 
               color = "black") +
  scale_color_manual(values=c('neg_sig'='blue','neg_insig'='skyblue','pos_sig'='red','pos_insig'='pink'))+
  scale_linetype_manual(values=c('insig.'=2,'sig.'=1))+
  geom_text(color="white", size=2) +
  labs(subtitle='Contexual',
       y='AME on suicide rates (per 100k)', x='') + 
  theme_pubr()+
  coord_flip() +
  theme(
    axis.text.y=element_text(size=8),
    axis.text.x=element_text(size=8),
    axis.title=element_text(size=8),
    plot.subtitle = element_text(hjust = 0.5))+
  guides(lty=FALSE,color=FALSE)

p = ggarrange(p1,p2, nrow=1)

ggsave(file=file.path(figure_path,file_name),plot=p, width=18,height=11,unit='cm',dpi=1000)

#------------------------------------------------------------------------------
# figures for sameness effects 
#==============================================================================
coef_for_plot = coef_combine[sameness=='sameness'&grepl('vs|Range',mtype),]
coef_for_plot = coef_for_plot[grepl('Range',mtype),]
coef_for_plot = coef_for_plot[at_type %in% c('Other factors','Religious Congregation')== FALSE,]

coef_for_plot = rbind(coef_for_plot[method=='logit',],
	coef_for_plot[method=='mi' & category %in% c('sameness UnEmpl','sameness PhysProb'),])

coef_for_plot[method=='logit',sig :=ifelse(p < 0.05, 'sig.','insig.')]
coef_for_plot[method=='mi',sig := ifelse(ll >0 | ul < 0, 'sig.','insig.')]

type_list = unique(coef_for_plot[,at_type])

coef_for_plot[,labels :=  round(change*100000)]

coef_for_plot[at_name=='% Same UnEmployed',at_name := '% Same Unemployed']

coef_for_plot[grep('Sex',at_name), category_order := 1]
coef_for_plot[grep('Age',at_name), category_order := 2]
coef_for_plot[grep('Marital Status',at_name), category_order := 3]
coef_for_plot[grep('Race',at_name), category_order := 4]
coef_for_plot[grep('BornUSA',at_name), category_order := 5]
coef_for_plot[grep('Unemployed',at_name), category_order := 6]
coef_for_plot[grep('Problem',at_name), category_order := 7]

coef_for_plot[change < 0 & sig == 'sig.',colors := 'neg_sig']
coef_for_plot[change < 0 & sig == 'insig.',colors := 'neg_insig']
coef_for_plot[change > 0 & sig == 'sig.',colors := 'pos_sig']
coef_for_plot[change > 0 & sig == 'insig.',colors := 'pos_insig']

file_name = 'figure2_lolly_pop_main_sameness_effects_county.pdf'

p = ggplot(coef_for_plot, aes(x=reorder(at_name,-category_order),
	y=change*100000,yend=ll*100000,ymax=ul*100000,label=round(change*100000,1),lty=sig,color=colors)) + 
  geom_point(stat='identity', size=10)  +
  geom_hline(yintercept=0,lty=2,lwd=0.1)+
  geom_segment(aes(y = ll*100000, 
                   x = at_name, 
                   yend = ul*100000, 
                   xend = at_name), 
               color = "black") +
  scale_color_manual(values=c('neg_sig'='blue','neg_insig'='skyblue','pos_sig'='red','pos_insig'='pink'))+
  scale_linetype_manual(values=c('insig.'=2,'sig.'=1))+
  geom_text(color="white", size=3) +
  labs(
       y='AME on suicide rates (per 100k)', x='') + 
  theme_pubr()+
  coord_flip() +
    theme(
    axis.text.y=element_text(size=8),
    axis.text.x=element_text(size=8),
    axis.title=element_text(size=8),
    plot.subtitle = element_text(hjust = 0.5))+

  guides(lty=FALSE,color=FALSE)

ggsave(file=file.path(figure_path,file_name),plot=p, width=11,height=8,unit='cm',dpi=1000)

#------------------------------------------------------------------------------
# comparison before/after individual/contextual controls 
#==============================================================================
# read all files 
coef_combine = list(); n = 1
	coef_combine[[n]] = read_output(log_path=log_path,sameness='raw',mm=1);n = n + 1
	coef_combine[[n]] = read_output(log_path=log_path,sameness='raw',mm=7);n = n + 1
	coef_combine[[n]] = read_output(log_path=log_path,sameness='raw',mm=8);n = n + 1
	coef_combine[[n]] = read_output_mi(log_path=log_path,sameness='raw',mm=2);n = n + 1
	coef_combine[[n]] = read_output_mi(log_path=log_path,sameness='raw',mm=9);n = n + 1
	coef_combine[[n]] = read_output_mi(log_path=log_path,sameness='raw',mm=10);n = n + 1
	
coef_combine = rbindlist(coef_combine,fill=TRUE)
#------------------------------------------------------------------------------
# main effets (raw)
#==============================================================================
coef_for_plot = coef_combine[sameness=='raw'&grepl('vs|Range',mtype),]
coef_for_plot = coef_for_plot[grep('vs 2|vs 3|vs 4',mtype,invert=TRUE),]
coef_for_plot = coef_for_plot[at_type %in% c('Other factors','Religious Congregation')== FALSE,]

coef_for_plot = rbind(coef_for_plot[method=='logit',],
	coef_for_plot[method=='mi' & category %in% c('UnEmpl','PhysProb','prop_UnEmpl','prop_PhysProb'),])

coef_for_plot[, group := ifelse(grepl('\\(vs ',category_name), 'individual','contextual')]
coef_for_plot[, group := factor(group, levels=c('individual','contextual'),
  labels=c('level: Individual','level: Contextual'))]

coef_for_plot[method=='logit',sig :=ifelse(p < 0.05, 'sig.','insig.')]
coef_for_plot[method=='mi',sig := ifelse(ll >0 | ul < 0, 'sig.','insig.')]

type_list = unique(coef_for_plot[,at_type])

coef_for_plot[,labels :=  round(change*100000)]

coef_for_plot[grep('Female',category_name), category_order := 1]
coef_for_plot[grep('Age',category_name), category_order := 2]
coef_for_plot[grep('Married|Widowed|Divorced|Separated|Never-married',category_name), category_order := 3]
coef_for_plot[grep('White|Black|Indian|Asian|Hispanic',category_name), category_order := 4]
coef_for_plot[grep('Native',category_name), category_order := 5]
coef_for_plot[grep('Unemployed',category_name), category_order := 6]
coef_for_plot[grep('Problem',category_name), category_order := 7]


coef_for_plot[change < 0 & sig == 'sig.',colors := 'neg_sig']
coef_for_plot[change < 0 & sig == 'insig.',colors := 'neg_insig']
coef_for_plot[change > 0 & sig == 'sig.',colors := 'pos_sig']
coef_for_plot[change > 0 & sig == 'insig.',colors := 'pos_insig']

coef_for_plot[,category_name:=gsub('\\[range\\]','',category_name)]
coef_for_plot[category_name =='Native (vson-native)',category_name:='Native (vs non-native)']


file_name = 'figure_lolly_pop_main_effects_county_beforeafter.pdf'

coef_for_plot[,model_type_label := factor(model_type, 
	levels=c(7,9,8,10,1,2),
  labels=c(
    'controls: individual only','controls: individual only',
		'controls: contextual only','controls: contextual only',
    'controls: both','controls: both'))]

p = ggplot(coef_for_plot, aes(x=reorder(category_name,-category_order),
	y=change*100000,yend=ll*100000,ymax=ul*100000,label=round(change*100000,1),lty=sig,color=colors)) + 
  geom_point(stat='identity', size=6)  +
  geom_hline(yintercept=0,lty=2,lwd=0.1)+
  geom_segment(aes(y = ll*100000, 
                   x = category_name, 
                   yend = ul*100000, 
                   xend = category_name), 
               color = "black") +
  scale_color_manual(values=c('neg_sig'='blue','neg_insig'='skyblue','pos_sig'='red','pos_insig'='pink'))+
  scale_linetype_manual(values=c('insig.'=2,'sig.'=1))+
  geom_text(color="white", size=2.5) +
  labs(title="Effects of individual and contextual factors on suicide rates", 
       subtitle="",
       y='Average marginal effects on suicide rates (per 100k)', x='') + 
  theme_pubr()+
  coord_flip() +
  theme(axis.text.y=element_text(size=8))+
  facet_wrap( ~ group+model_type_label, scale='free') +
  guides(lty=FALSE,color=FALSE)

ggsave(file=file.path(figure_path,file_name),plot=p, width=8,height=8,dpi=1000)

#------------------------------------------------------------------------------
# interaction effects
#==============================================================================
outcome = 'sameness'

target_category = list()
target_category[[1]] = list(name='Female',value=c(0,1),value_name=c('Male','Female'))
target_category[[2]] = list(name='AgeGrp4',value=c(1,2,3,4),value_name=c('15_24','25_44','45_64','65_Up'))
target_category[[3]] = list(name='Race5',value=c(1,2,3,4,5),value_name=c('White','Black','Indian|Alaska','Asian','Hispanic'))
target_category[[4]] = list(name='MarStat5',value=c(1,2,3,4,5),value_name=c('Married','Widowed','Divorced','Separated','Never Married'))
target_category[[5]] = list(name='BornUSA',value=c(0,1),value_name=c('not US-born','US-born'))
target_category[[6]] = list(name='UnEmpl',value=c(0,1),value_name=c('Employed','Unemployed'))
target_category[[7]] = list(name='PhysProb',value=c(0,1),value_name=c('Healthy','Physical Problem'))

target_labels = c('A. Age','B. Race','C. Nativity',
    'D. Marital Status','E. Physical Problem','F. Unemployment')

list_p = list(); n = 1
for (i in c(2,3,5,4,7,6)){
	target_var = target_category[[i]]$name
	if (i <= 5) {
		model_no = 5
		method ='logit'
	} else {
		model_no = 6 
		method = 'mi'
	}
	
  p <- plot_interaction(method=method,log_path=log_path,outcome = outcome,target_var = target_var,model_no=model_no)
  p = p + theme(
    axis.text.y=element_text(size=8),
    axis.text.x=element_text(size=8),
    axis.title.y=element_text(size=10),
    axis.title.x=element_text(size=7),
    plot.subtitle = element_text(hjust = 0,size=10,face='bold'))
  p = p + labs(subtitle=target_labels[n])
  list_p[[n]] = p 
  n = n + 1
}

p = ggarrange(
  list_p[[1]],list_p[[2]],list_p[[3]],
  list_p[[4]],list_p[[5]],list_p[[6]],
  ncol=3,nrow=2)

file_name = 'figure3_sameness_interaction_effects.pdf'
ggsave(file=file.path(figure_path,file_name),plot=p, width=18,height=15,unit='cm',dpi=1000)
